"""
yfinance数据源适配器
提供与现有数据源兼容的接口，作为备用数据源
"""

import yfinance as yf
import time
import logging
from typing import Dict, List, Optional
import pandas as pd
from concurrent.futures import ThreadPoolExecutor, as_completed


class YFinanceAdapter:
    """yfinance数据源适配器"""
    
    def __init__(self):
        self.logger = self._setup_logging()
        self.session = None
        self._setup_session()
    
    def _setup_logging(self) -> logging.Logger:
        """设置日志"""
        logger = logging.getLogger('YFinanceAdapter')
        if not logger.handlers:
            handler = logging.StreamHandler()
            formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
            handler.setFormatter(formatter)
            logger.addHandler(handler)
            logger.setLevel(logging.INFO)
        return logger
    
    def _setup_session(self):
        """设置会话，减少连接开销"""
        try:
            # yfinance内部会处理session，这里主要是为了兼容性
            self.session = True
            self.logger.info("YFinance adapter initialized")
        except Exception as e:
            self.logger.error(f"Failed to setup YFinance session: {e}")
    
    def _convert_stock_code(self, code: str) -> str:
        """将A股代码转换为Yahoo Finance格式
        
        Args:
            code: 6位A股代码
            
        Returns:
            Yahoo Finance格式的股票代码
        """
        if not code or len(code) != 6:
            return code
        
        # A股代码转换规则
        if code.startswith('6'):
            # 沪市主板：600xxx -> 600xxx.SS
            return f"{code}.SS"
        elif code.startswith('000'):
            # 深市主板：000xxx -> 000xxx.SZ
            return f"{code}.SZ"
        elif code.startswith('002'):
            # 中小板：002xxx -> 002xxx.SZ
            return f"{code}.SZ"
        elif code.startswith('300'):
            # 创业板：300xxx -> 300xxx.SZ
            return f"{code}.SZ"
        elif code.startswith('688'):
            # 科创板：688xxx -> 688xxx.SS
            return f"{code}.SS"
        else:
            # 其他情况，默认深市
            return f"{code}.SZ"
    
    def _fetch_single_stock(self, code: str, max_retries: int = 2) -> Optional[Dict]:
        """获取单只股票数据
        
        Args:
            code: 股票代码
            max_retries: 最大重试次数
            
        Returns:
            股票数据字典或None
        """
        yf_code = self._convert_stock_code(code)
        
        for attempt in range(max_retries + 1):
            try:
                ticker = yf.Ticker(yf_code)
                
                # 获取实时数据
                info = ticker.info
                hist = ticker.history(period="2d")
                
                if info and not hist.empty:
                    # 转换为兼容格式
                    current_price = hist['Close'].iloc[-1] if len(hist) > 0 else info.get('currentPrice', 0)
                    previous_close = hist['Close'].iloc[-2] if len(hist) > 1 else info.get('previousClose', current_price)
                    
                    # 如果没有历史数据，尝试从info获取
                    if current_price == 0:
                        current_price = info.get('currentPrice', 0)
                    if previous_close == 0:
                        previous_close = info.get('previousClose', current_price)
                    
                    stock_data = {
                        'name': info.get('longName', info.get('shortName', f'股票{code}')),
                        'now': float(current_price) if current_price else 0.0,
                        'close': float(previous_close) if previous_close else 0.0,
                        'high': float(hist['High'].iloc[-1]) if not hist.empty else float(current_price),
                        'low': float(hist['Low'].iloc[-1]) if not hist.empty else float(current_price),
                        'open': float(hist['Open'].iloc[-1]) if not hist.empty else float(current_price),
                        'turnover': int(hist['Volume'].iloc[-1]) if not hist.empty else 0,
                        'volume': int(hist['Volume'].iloc[-1]) if not hist.empty else 0,
                    }
                    
                    # 验证数据有效性
                    if stock_data['now'] > 0:
                        return stock_data
                    else:
                        self.logger.warning(f"Invalid price data for {code}: {stock_data['now']}")
                        
                else:
                    self.logger.warning(f"No data available for {yf_code}")
                    
            except Exception as e:
                if attempt < max_retries:
                    self.logger.warning(f"Attempt {attempt + 1} failed for {code}: {e}, retrying...")
                    time.sleep(1)  # 短暂延迟后重试
                else:
                    self.logger.error(f"Failed to fetch data for {code} after {max_retries + 1} attempts: {e}")
        
        return None
    
    def stocks(self, codes: List[str], max_workers: int = 5) -> Dict[str, Dict]:
        """批量获取股票数据
        
        Args:
            codes: 股票代码列表
            max_workers: 最大并发数
            
        Returns:
            股票数据字典
        """
        if not codes:
            return {}
        
        self.logger.info(f"Fetching data for {len(codes)} stocks using YFinance")
        
        results = {}
        
        # 使用线程池并发获取数据，但限制并发数避免速率限制
        with ThreadPoolExecutor(max_workers=min(max_workers, len(codes))) as executor:
            # 提交任务
            future_to_code = {
                executor.submit(self._fetch_single_stock, code): code 
                for code in codes
            }
            
            # 收集结果
            for future in as_completed(future_to_code):
                code = future_to_code[future]
                try:
                    data = future.result(timeout=10)  # 10秒超时
                    if data:
                        results[code] = data
                    else:
                        self.logger.warning(f"No data returned for {code}")
                except Exception as e:
                    self.logger.error(f"Error processing {code}: {e}")
                
                # 添加小延迟避免速率限制
                time.sleep(0.1)
        
        success_rate = len(results) / len(codes) * 100 if codes else 0
        self.logger.info(f"YFinance fetch completed: {len(results)}/{len(codes)} ({success_rate:.1f}%)")
        
        return results
    
    def test_connection(self) -> bool:
        """测试连接
        
        Returns:
            连接是否成功
        """
        try:
            # 测试获取一个知名股票的数据
            test_data = self._fetch_single_stock('000001')  # 平安银行
            return test_data is not None
        except Exception as e:
            self.logger.error(f"Connection test failed: {e}")
            return False


class YFinanceQuotation:
    """yfinance行情服务，兼容easyquotation接口"""
    
    def __init__(self):
        self.adapter = YFinanceAdapter()
    
    def stocks(self, codes: List[str]) -> Dict[str, Dict]:
        """获取股票数据，兼容easyquotation接口"""
        return self.adapter.stocks(codes)
    
    def test_connection(self) -> bool:
        """测试连接"""
        return self.adapter.test_connection()


def create_yfinance_quotation() -> YFinanceQuotation:
    """创建yfinance行情服务实例"""
    return YFinanceQuotation() 